%% Replication File for Baqaee, Burstein, and Koike-Mori (2023)  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Prepared by Yasutaka Koike-Mori (UCLA) 
% Description: This function calculates the Money Metric Utility Function allowing missing price using initial-year base prices.
% The order for all matrices is (I,N,T) where:
% I denotes the number of goods
% N denotes the number of households.
% T denotes the number of time periods
% 
% Inputs: I_vec, Bx_vec, Bxi_vec, pvec
%   I_vec:     A (1,N,T) matrix representing income.
%   Bx_vec   : A (1,N,T) matrix of Uncompensated Budget share of observed goods (total)
%   Bxi_vec  : A (I,N,T) matrix of Uncompensated Budget shares within the observed group
%   pvec :     A (1,N,T) matrix representing price.
%   sigmaM: Marshallian Elasticity (A scalar or (1,N,T) matrix)
%           A scalar calibrated value can be used, or a matrix if it is a functions of Income.
%                              
%   flagAlgorithm: Set to 0 for iterative, 1 for recursive.
%   Note: Prior to use, sort I_vec and B_vec in ascending order by the second dimension of I_vec.
%         The number of households can vary by period. See N to be the maximum number of households across periods, and set data for missing households in a give period to NaN
%
% Outputs: U_vec, WP_total_vec and WP_FS_vec
%   U_vec       : A (1,N,T) matrix containing Money Metric. For income levels for which we cannot calculate money-metric, U_vec=NaN.
%   WP_total_vec: Inflation (for decomposition figure)
%   WP_FS_vec   : Missing price term (for decomposition figure)
%
% Internal Functions: compensated_bx, compensated_bxi, compensated_elast, CalMoneyMetric_t, updateMMindex, makefunction
%   uncompensated_Bx, uncompensated_Bxi and  Compensated_elast: Calculates compensated functions.
%   CalMoneyMetric_t   : Calculates money metric at time t.
%   checkIterative     : Final check for iterative algorithm.
%   makefunction       : Interpolates allowing for NaN values (Allow approximations outside the support).
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [U_vec, WP_total_vec, WP_FS_vec] = CalMoneyMetric_Unobserved(I_vec, Bx_vec, Bxi_vec, pvec, sigmaM, flagAlgorithm)
    [~, N, T] = size(Bxi_vec);
    WP_total_vec = zeros(1, N, T);
    WP_FS_vec = zeros(1, N, T);
    U_vec(:,:,1) = I_vec(:,:,1); 
    if flagAlgorithm == 1 % check if iterative or recursive
        Maxiter = 20;
    else
        Maxiter = 1;
    end
    % Initializion
    dlogp = [diff(log(pvec),[],3)];
    U_vec(:,:,1) = I_vec(:,:,1);
    F_Bxi_I = uncompensated_Bxi(Bxi_vec,I_vec,T); F_Bx_I= uncompensated_Bx(Bx_vec,I_vec,T); F_sigmaI= uncompensated_elast(sigmaM,I_vec,T);  % uncompensated functions of Income without extrapolation
    F_I_U  = cell(1,T);  F_U_I  = cell(1,T);

    % Loop over T
    for t = 2:T    
        if mod(t,10) == 0
            disp(['t: ', num2str(t)]);
        end
        % Initialization
        tau = 1; err = 100; tol = 10^(-4); logU_vec_old = log(I_vec(:,:,t)); 
        % Iterate until convergence
        while err > tol
            [U_vec,WP_total_vec, WP_FS_vec,F_I_U,F_U_I] = CalMoneyMetric_t(t, dlogp, I_vec, U_vec, Bxi_vec, Bx_vec, WP_total_vec, WP_FS_vec,F_Bxi_I,F_Bx_I,F_sigmaI ,F_I_U,F_U_I,tau);
            if tau > 1
                err = max(abs(log(U_vec(:,:,t)) - logU_vec_old));
            end   
            logU_vec_old = log(U_vec(:,:,t));            
            tau = tau + 1;
            if tau > Maxiter
                break
            end
        end
    end 
end


function [U_vec,WP_total_vec, WP_FS_vec,F_I_U,F_U_I] = CalMoneyMetric_t(t, dlogp, I_vec, U_vec, Bxi_vec, Bx_vec, WP_total_vec, WP_FS_vec,F_Bxi_I,F_Bx_I,F_sigmaI ,F_I_U,F_U_I,tau)
    % Get sizes
    [I, N, ~] = size(Bxi_vec);
    % Initialization
    Integ = zeros(1, N);
    Integ_FS = zeros(1, N);
    Bxi_Istar = zeros(I,N);
    % Compute utility values for time step t
    if tau ==1
        % Make function of I(U) and U(I). This function allows approximation outside of support
        F_I_U{1,t-1} = makefunction(log(U_vec(:,:,t-1)),log(I_vec(:,:,t-1)));
        F_U_I{1,t-1} = makefunction(log(I_vec(:,:,t-1)),log(U_vec(:,:,t-1)));   
        % Initial Guess:
        logUt = F_U_I{1,t-1}(log(I_vec(:,:,t)));
    else % If tau > 1, use the value of previous iteration.
        logUt = log(U_vec(:,:,t));  
    end 

    for s = 1:t-1
        logIstar_s = F_I_U{1,s}(logUt);
        Bx_Istar_s  = F_Bx_I{1,s}(logIstar_s);
        if s < t-1
            logIstar_s1 = F_I_U{1,s+1}(logUt); 
            Bx_Istar_s1 = F_Bx_I{1,s+1}(logIstar_s1);
        else
            Bx_Istar_s1 = Bx_vec(:,:,s+1);
        end                     

        for i = 1:I          
            if s < t-1 
                Bxi_Istar(i,:) = (F_Bxi_I{i,s}(logIstar_s) + F_Bxi_I{i,s+1}(logIstar_s1))/2; 
            else
                Bxi_Istar(i,:) = (F_Bxi_I{i,s}(logIstar_s) + Bxi_vec(i,s+1))/2;
            end
        end 
        Bxi_Istar = Bxi_Istar./sum(Bxi_Istar,1); 
        sigmaIstar = F_sigmaI{1,s}(logIstar_s); 

        dlogbx = log(Bx_Istar_s1) - log(Bx_Istar_s);
        Integ = real(Integ + ( - (sum(Bxi_Istar.*dlogp(:,:,s))) - 1./(sigmaIstar -1).*dlogbx )); 
        Integ_FS = Integ_FS - 1./(sigmaIstar -1).*dlogbx; 
    end
    % Update U_vec
    % Since F_B_I returns NaN if the income is outside the support of the Engel Curve, U_vec contains only samples without extrapolation of Engel Curve.
    U_vec(:,:,t) = exp(log(I_vec(:,:,t)) + Integ);
    % Final check for iterative solution
    if tau == 1
        U_vec = checkIterative(U_vec,I_vec,F_I_U, t);
    end
    % save for the figures
    WP_total_vec(:,:,t) = Integ;
    WP_FS_vec(:,:,t) = Integ_FS;
end

%% Internal functions
% Final check for iterative solution
function U_vec = checkIterative(U_vec,I_vec,F_I_U, t)
for s = 1:t-1
    maxIs = max(I_vec(:,:,s));
    minIs = min(I_vec(:,:,s));
    Istars = exp(F_I_U{1,s}(log(U_vec(:,:,t))));
    indexs = (Istars >= minIs) & (Istars <= maxIs) ;
    U_vec(:,indexs==0,t) = NaN;
end
end

% unconpensated budget share (Bxi)
function f_Bxi_I = uncompensated_Bxi(Bxi_vec,I_vec,T)
    for t = 1:T
        for i = 1:size(Bxi_vec,1)
            f_Bxi_I{i,t} = interp_I(log(I_vec(:,:,t)),Bxi_vec(i,:,t));
        end
    end
end

% unconpensated budget share (Bx)
function f_Bx_I = uncompensated_Bx(Bx_vec,I_vec,T)
   for t = 1:T
      f_Bx_I{1,t} = interp_I(log(I_vec(:,:,t)), Bx_vec(:,:,t));
   end
end

% unconpensated elasticity
function f_sigmaI = uncompensated_elast(sigmaM,I_vec,T)
    if isscalar(sigmaM)
        sigmaM = sigmaM.*ones(size(I_vec));
    end
    for t = 1:T
       f_sigmaI{1,t} = interp_I(log(I_vec(:,:,t)),sigmaM(:,:,t));
    end
end


% Function to interpolate vectors with NaN values
function F_X = makefunction(X,Y)
    indexToKeep = ~isnan(X) & ~isnan(Y);
    X = X(indexToKeep); Y = Y(indexToKeep);    
    [X, indexUN] = unique(X);
    Y = Y(indexUN);
    F_X = griddedInterpolant(X,Y);
end

% Function to interpolate a function of I without extrapolation
function F_X = interp_I(X,Y)
    indexToKeep = ~isnan(X) & ~isnan(Y);
    X = X(indexToKeep); Y = Y(indexToKeep);
    [X, indexUN] = unique(X);
    Y = Y(indexUN);
    F_X = griddedInterpolant(X,Y,'linear','none');
end


